package com.lihh;

import org.apache.zookeeper.*;
import org.apache.zookeeper.data.Stat;

import java.io.IOException;
import java.util.HashSet;
import java.util.List;
import java.util.Set;

/**
 * zk 的客户端工具
 *
 * @author lihh
 */
public class ZkClientUtils {
  // 表示连接的客户端
  private static ZooKeeper client;
  // 表示注册的服务地址
  private static Set<String> servers = new HashSet<>();
  // 表示 root 根节点
  private static String ROOT_HOST_NAME = "/servers";
  // 集群地址
  private final static String connectString = "192.168.56.17:2181,192.168.56.18:2181,192.168.56.19:2181";
  
  /**
   * 拿到注册的服务地址
   *
   * @return 返回服务地址列表
   * @author lihh
   */
  
  public static Set<String> getRegisterServers() {
    return servers;
  }
  
  /**
   * 连接 地址
   *
   * @author lihh
   */
  public static void connect() throws IOException {
    client = new ZooKeeper(connectString, 10000, new Watcher() {
      @Override
      public void process(WatchedEvent watchedEvent) {
        try {
          getChildrenList();
        } catch (InterruptedException e) {
          throw new RuntimeException(e);
        } catch (KeeperException e) {
          throw new RuntimeException(e);
        }
      }
    });
  }
  
  /**
   * 拿到 /servers 子节点
   *
   * @author lihh
   */
  public static void getChildrenList() throws InterruptedException, KeeperException {
    // 判断节点是否存在
    Stat stat = client.exists(ROOT_HOST_NAME, false);
    if (stat == null) return;
    
    List<String> children = client.getChildren(ROOT_HOST_NAME, true);
    children.forEach(node -> {
      try {
        byte[] data = client.getData(ROOT_HOST_NAME + "/" + node, false, null);
        servers.add(new String(data));
      } catch (KeeperException e) {
        throw new RuntimeException(e);
      } catch (InterruptedException e) {
        throw new RuntimeException(e);
      }
      
    });
  }
  
  /**
   * 添加根节点
   *
   * @throws InterruptedException
   * @throws KeeperException
   * @author lihh
   */
  private static void createRootNode() throws InterruptedException, KeeperException {
    // 判断节点是否存在
    Stat stat = client.exists(ROOT_HOST_NAME, false);
    // 如果节点不存的话 && 创建节点
    if (stat == null)
      client.create(ROOT_HOST_NAME, ROOT_HOST_NAME.getBytes(), ZooDefs.Ids.OPEN_ACL_UNSAFE, CreateMode.PERSISTENT);
  }
  
  /**
   * 进行注册的内容
   *
   * @param path 添加path 节点
   * @author lihh
   */
  public static void register(String path) throws InterruptedException, KeeperException {
    if (!path.startsWith("/"))
      new Exception("zookeeper node[path] must start with[/]");
    
    createRootNode();
    // 开始创建虚拟节点
    String createNodePath = ROOT_HOST_NAME + path;
    client.create(createNodePath, path.substring(1).getBytes(), ZooDefs.Ids.OPEN_ACL_UNSAFE, CreateMode.EPHEMERAL_SEQUENTIAL);
    
    servers.add(path.substring(1));
  }
  
  /**
   * 关闭连接
   *
   * @author lihh
   */
  public static void close() throws InterruptedException {
    client.close();
  }
}
